{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Problem: Add a Benchmark to Your PyTorch Code\n",
    "\n",
    "### Problem Statement\n",
    "You are tasked with implementing a simple neural network model with fully connected layers and adding benchmarking functionality to measure and display the time taken for each epoch of training and testing. The goal is to evaluate the model's performance and record the time taken for both training and testing phases.\n",
    "\n",
    "### Requirements\n",
    "1. **Define a Neural Network Model**:\n",
    "   - Implement a simple feedforward neural network using fully connected layers (`nn.Linear`).\n",
    "   - The network should be suitable for classification tasks.\n",
    "\n",
    "2. **Benchmark Training and Testing**:\n",
    "   - Measure the time taken for each epoch during training and display the elapsed time.\n",
    "   - Measure and display the time taken for the testing phase after each epoch.\n",
    "\n",
    "### Constraints\n",
    "- The model should have at least two hidden layers with ReLU activations.\n",
    "- Use the appropriate loss function and optimizer for training the model.\n",
    "- Ensure that the benchmarking measures both the training and testing time accurately.\n",
    "\n",
    "<details>\n",
    "  <summary>💡 Hint</summary>\n",
    "  Define the SimpleNN class:\n",
    "  <br>\n",
    "  Add two fully connected layers:\n",
    "  <br>\n",
    "  Apply a ReLU activation function to the first layer.\n",
    "  <br>\n",
    "  <br>\n",
    "  Benchmark the Code:\n",
    "  <br>\n",
    "  Measure and print training time for each epoch.\n",
    "  <br>\n",
    "  Measure and print testing time along with accuracy.\n",
    "</details>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import torchvision\n",
    "import torchvision.transforms as transforms\n",
    "import time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load MNIST dataset\n",
    "transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])\n",
    "\n",
    "train_dataset = torchvision.datasets.MNIST(root='./data', train=True, download=True, transform=transform)\n",
    "train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=64, shuffle=True)\n",
    "\n",
    "test_dataset = torchvision.datasets.MNIST(root='./data', train=False, download=True, transform=transform)\n",
    "test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=64, shuffle=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define a simple neural network model\n",
    "# TODO: Add layers to the model\n",
    "class SimpleNN(nn.Module):\n",
    "    def __init__(self):\n",
    "        ...\n",
    "\n",
    "    def forward(self, x):\n",
    "        ...\n",
    "\n",
    "# Initialize the model, loss function, and optimizer\n",
    "model = SimpleNN()\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = optim.SGD(model.parameters(), lr=0.01)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch [1/5], Loss: 0.3573, Time: 8.3633s\n",
      "Epoch [2/5], Loss: 0.4543, Time: 7.1758s\n",
      "Epoch [3/5], Loss: 0.3774, Time: 20.4998s\n",
      "Epoch [4/5], Loss: 0.2720, Time: 31.5872s\n",
      "Epoch [5/5], Loss: 0.2737, Time: 7.2157s\n"
     ]
    }
   ],
   "source": [
    "# Training loop with benchmarking\n",
    "epochs = 5\n",
    "for epoch in range(epochs):\n",
    "    start_time = time.time()  # Start time for training\n",
    "    for images, labels in train_loader:\n",
    "        # Forward pass\n",
    "        outputs = model(images)\n",
    "        loss = criterion(outputs, labels)\n",
    "\n",
    "        # Backward pass and optimization\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "    end_time = time.time()  # End time for training\n",
    "    training_time = end_time - start_time\n",
    "    print(f\"Epoch [{epoch + 1}/{epochs}], Loss: {loss.item():.4f}, Time: {training_time:.4f}s\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test Accuracy: 92.71%, Testing Time: 0.9687s\n"
     ]
    }
   ],
   "source": [
    "# Evaluate the model on the test set and benchmark the accuracy\n",
    "correct = 0\n",
    "total = 0\n",
    "start_time = time.time()  # Start time for testing\n",
    "with torch.no_grad():\n",
    "    for images, labels in test_loader:\n",
    "        outputs = model(images)\n",
    "        _, predicted = torch.max(outputs, 1)\n",
    "        total += labels.size(0)\n",
    "        correct += (predicted == labels).sum().item()\n",
    "\n",
    "end_time = time.time()  # End time for testing\n",
    "testing_time = end_time - start_time\n",
    "accuracy = 100 * correct / total\n",
    "print(f\"Test Accuracy: {accuracy:.2f}%, Testing Time: {testing_time:.4f}s\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
